from __future__ import print_function, division

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.autograd import Variable
from torchvision import models
import cv2
import numpy as np
import os
from lib.network import Network
import torchvision


def del_dir(path):
    for i in os.listdir(path):
        path_file = os.path.join(path, i)       # 取文件绝对路径
        if os.path.isfile(path_file):
            os.remove(path_file)
        else:
            del_dir(path_file)


def load_model(test_dir):
    # test_datasets = torchvision.datasets.MNIST(root='./mnist/',
    #                                        transform=torchvision.transforms.ToTensor(),
    #                                        train=False)
    val_transforms = transforms.Compose([
        transforms.Grayscale(),
        transforms.ToTensor()
    ])
    test_datasets = datasets.ImageFolder(test_dir, transform=val_transforms)

    net = Network()
    if torch.cuda.is_available():
        net = nn.DataParallel(net)
        net.cuda()
    state_dict = torch.load(r".\mnist_net.pth")
    net.load_state_dict(state_dict)
    return net, test_datasets


def Tensor2Cvimg(img_tensor):
    img_np = img_tensor.numpy()
    img_np = np.transpose(img_np, (1, 2, 0))
    img_np = (img_np * 0.5 + 0.5) * 255
    img_cv = img_np.astype("uint8")
    img_cv = cv2.cvtColor(img_cv, cv2.COLOR_RGB2BGR)
    return img_cv


def Cvimg2Tensor(img_cv):
    img_cv = cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB)
    img_tensor = torch.from_numpy(np.transpose(img_cv, (2, 0, 1)))
    return img_tensor


def imread_tensor(img_path):
    img_cv = cv2.imread(img_path, 1)
    img_tensor = Cvimg2Tensor(img_cv)
    return img_tensor


def run():
    test_dir = r'H:\yuanbaoxi\ybx_gitee\non_local\mnist\testimgs'
    net, test_datasets = load_model(test_dir)
    print("begin predict....")
    del_dir("./result")
    with torch.no_grad():
        cls = test_datasets.classes
        class_num = len(cls)
        img_names = test_datasets.imgs
        img_num = test_datasets.__len__()
        for i in range(0, img_num):
            net.eval()
            img, target = test_datasets.__getitem__(i)
            size = img.size()
            img = img.view(1,size[0], size[1], size[2])
            t_target = torch.from_numpy(np.array(target))
            inputs, labels = Variable(img).cuda(), Variable(t_target).cuda()
            outputs = net(inputs)
            # 取得分最高的那个类 (outputs.data的索引号)
            _, predicted = torch.max(outputs.data, 1)
            predicted_cpu = Variable(predicted).cpu()
            predicted_np = predicted_cpu.numpy()
            pred = int(predicted_np[0])
            lbl = int(target)
            print("pred result %s\n"%cls[pred])
            save_path = "./temp.jpg"
            img = img.view(size[0], size[1], size[2])
            img_cv = Tensor2Cvimg(img)
            cv2.putText(img_cv, '%s' % cls[pred], (0, 20), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 1)
            cv2.imwrite(save_path, img_cv)
            cv2.imshow("result", img_cv)
            cv2.waitKey(2000)


if __name__ == "__main__":
    run()